##########################################################################################
# Title: SA-Albers (single CRS) Buildings
# Author: Varun Kshirsagar 
# Date: 2025-10-01
# Cities: Cape Town, Johannesburg, Durban (eThekwini)
# Approach:
#   • WGS84 (EPSG:4326) only for bbox + initial city clip (s2 geodesic on)
#   • One projected CRS for everything else: South Africa Albers (equal-area)
#   • Robust geometry handling to avoid st_cast errors on messy subplace/building geoms
##########################################################################################

suppressPackageStartupMessages({
  library(sf)
  library(sfarrow)
  library(arrow)
  library(future)
  library(furrr)
  library(dplyr)
  library(data.table)
  library(units)
  library(lwgeom)
})

## --------- Tunables ----------
CORES                   <- min(5, max(1, parallel::detectCores() - 3))
BATCH_SIZE              <- 25000L      # ↓ if RAM tight
CELL_M                  <- 200         # grid cell size (meters)
SIMPLIFY_BOUNDARY_M     <- 5           # boundaries only
SIMPLIFY_BUILDINGS_M    <- 0           # keep buildings intact
WRITE_BUILDINGS_GPKG    <- TRUE

options(future.globals.maxSize = 8 * 1024^3)
options(future.rng.onMisuse = "ignore")
plan(multisession, workers = CORES)

## --------- CRS ----------
CRS_WGS84     <- sf::st_crs(4326)
CRS_SA_ALBERS <- sf::st_crs(
  "+proj=aea +lat_1=-18 +lat_2=-32 +lat_0=0 +lon_0=25 +datum=WGS84 +units=m +no_defs"
)

## --------- Paths ----------
if (!requireNamespace("here", quietly = TRUE)) install.packages("here", quiet = TRUE)
ROOT <- normalizePath(here::here(), winslash = "/")
PARQUET_DIR   <- file.path(ROOT, "Data", "GOB")
BOUNDARY_DIR  <- file.path(ROOT, "Data", "Shape_Files")
SUBPLACE_FILE <- file.path(ROOT, "Data", "Shape_Files", "SP_SA_2011_clean.gpkg")
OUT_DIR       <- file.path(ROOT, "Processed_Data", "Buildings")
dir.create(OUT_DIR, showWarnings = FALSE, recursive = TRUE)
stopifnot(file.exists(SUBPLACE_FILE))

## --------- City configs (NO autodetect) ----------
city_cfg <- list(
  capetown = list(
    city_id       = "capetown",
    parquet_file  = file.path(PARQUET_DIR, "2148217022255726592.parquet"),
    boundary_file = file.path(BOUNDARY_DIR, "capetown_boundary_sp2011_clean.gpkg"),
    admin_keep    = "City of Cape Town"
  ),
  johannesburg = list(
    city_id       = "johannesburg",
    parquet_file  = file.path(PARQUET_DIR, "2202260217784172544.parquet"),
    boundary_file = file.path(BOUNDARY_DIR, "johannesburg_boundary_sp2011_clean.gpkg"),
    admin_keep    = "City of Johannesburg"
  ),
  durban = list(
    city_id       = "durban",
    parquet_file  = file.path(PARQUET_DIR, "2229281815548395520.parquet"),
    boundary_file = file.path(BOUNDARY_DIR, "durban_boundary_sp2011_clean.gpkg"),
    admin_keep    = "eThekwini"
  )
)

## --------- Helpers ----------
try_suppress <- function(expr) { suppressWarnings(try(expr, silent = TRUE)) }

# Boundary from prebuilt city boundary GPKG
build_city_boundary <- function(admin_keep, boundary_file) {
  b_raw <- sf::st_read(boundary_file, quiet = TRUE) |> sf::st_make_valid()
  if (nrow(b_raw) > 1) b_raw <- sf::st_union(b_raw)
  if (is.na(sf::st_crs(b_raw))) sf::st_crs(b_raw) <- CRS_WGS84
  b_wgs <- sf::st_transform(b_raw, CRS_WGS84)
  b_sa  <- sf::st_transform(b_wgs, CRS_SA_ALBERS) |>
    sf::st_make_valid() |>
    sf::st_set_precision(0.01) |>
    sf::st_simplify(dTolerance = SIMPLIFY_BOUNDARY_M)
  list(
    BOUND_WGS84      = b_wgs,
    BOUND_WGS84_GEOM = sf::st_geometry(b_wgs),
    BOUND_BBOX_SFC   = sf::st_as_sfc(sf::st_bbox(b_wgs)),
    BOUND_SA         = b_sa,
    BOUND_SA_GEOM    = sf::st_geometry(b_sa)
  )
}

# Chunk processor: bbox + boundary clip in WGS84; write WGS84 chunks
make_process_chunk_wgs84 <- function(BOUND_BBOX_SFC, BOUND_WGS84_GEOM, TMP_DIR) {
  force(BOUND_BBOX_SFC); force(BOUND_WGS84_GEOM); force(TMP_DIR)
  function(ix, chunk_sf) {
    if (is.na(st_crs(chunk_sf))) st_crs(chunk_sf) <- CRS_WGS84
    if (st_crs(chunk_sf)$epsg != 4326) chunk_sf <- st_transform(chunk_sf, CRS_WGS84)
    
    # Fast bbox filter
    idx_bbox  <- st_intersects(chunk_sf, BOUND_BBOX_SFC, sparse = TRUE, prepared = TRUE)
    keep_bbox <- lengths(idx_bbox) > 0L
    if (!any(keep_bbox)) return(NA_character_)
    chunk_sf <- chunk_sf[keep_bbox, , drop = FALSE]
    
    # Boundary clip (geodesic)
    old_s2 <- sf::sf_use_s2(); on.exit(sf::sf_use_s2(old_s2), add = TRUE)
    sf::sf_use_s2(TRUE)
    
    c_try <- try_suppress(st_intersection(st_make_valid(chunk_sf), BOUND_WGS84_GEOM))
    if (inherits(c_try, "try-error") || is.null(c_try) || nrow(c_try) == 0) {
      idx_real  <- st_intersects(chunk_sf, BOUND_WGS84_GEOM, sparse = TRUE, prepared = TRUE)
      keep_real <- lengths(idx_real) > 0L
      if (!any(keep_real)) return(NA_character_)
      c_try <- chunk_sf[keep_real, , drop = FALSE]
    }
    
    # Robust polygonization: drop non-polygonal parts before casting
    c_try <- st_make_valid(c_try)
    c_try <- c_try[!st_is_empty(c_try), , drop = FALSE]
    if (!nrow(c_try)) return(NA_character_)
    
    # Extract polygonal only; if nothing polygonal remains, skip
    c_poly <- st_collection_extract(c_try, type = "POLYGON", warn = FALSE)
    if (!nrow(c_poly) || all(st_is_empty(c_poly))) return(NA_character_)
    
    # Safe cast
    c_poly <- st_cast(c_poly, "MULTIPOLYGON", warn = FALSE)
    
    chunk_path <- file.path(TMP_DIR, sprintf("buildings_chunk_%04d.gpkg", ix))
    st_write(c_poly, chunk_path, delete_dsn = TRUE, quiet = TRUE)
    return(chunk_path)
  }
}

merge_chunk_files_wgs84 <- function(chunk_files, out_gpkg) {
  if (length(chunk_files) == 0) return(invisible(FALSE))
  if (file.exists(out_gpkg)) file.remove(out_gpkg)
  for (k in seq_along(chunk_files)) {
    ch <- st_read(chunk_files[k], quiet = TRUE)
    if (is.na(st_crs(ch))) st_crs(ch) <- CRS_WGS84
    if (st_crs(ch)$epsg != 4326) ch <- st_transform(ch, CRS_WGS84)
    ch <- st_collection_extract(ch, type = "POLYGON", warn = FALSE)
    ch <- ch[!st_is_empty(ch), , drop = FALSE]
    if (!nrow(ch)) next
    ch <- st_cast(ch, "MULTIPOLYGON", warn = FALSE)
    st_write(ch, out_gpkg, append = file.exists(out_gpkg), quiet = TRUE)
  }
  invisible(TRUE)
}

# Assign buildings to grids via centroid; aggregate counts/areas
grid_stats_from_buildings <- function(buildings_sa, grid_sf) {
  stopifnot("grid_id" %in% names(grid_sf))
  buildings_sa <- sf::st_make_valid(buildings_sa)
  buildings_sa <- buildings_sa[!sf::st_is_empty(buildings_sa), ]
  if (nrow(buildings_sa) == 0) {
    return(dplyr::mutate(grid_sf, n_buildings = 0L, total_area_m2 = 0, avg_build_size = 0))
  }
  b_pts <- sf::st_centroid(buildings_sa)
  if (sf::st_crs(b_pts) != sf::st_crs(grid_sf)) b_pts <- sf::st_transform(b_pts, sf::st_crs(grid_sf))
  b_pts <- sf::st_sf(bid = seq_len(nrow(b_pts)), geometry = sf::st_geometry(b_pts), crs = sf::st_crs(grid_sf))
  j <- sf::st_join(b_pts, grid_sf["grid_id"], left = FALSE)
  if (nrow(j) == 0) {
    return(dplyr::mutate(grid_sf, n_buildings = 0L, total_area_m2 = 0, avg_build_size = 0))
  }
  dt <- data.table::data.table(bid = j$bid, grid_id = j$grid_id, area_m2 = buildings_sa$area_m2[j$bid])
  grid_stats <- dt[, .(n_buildings = .N, total_area_m2 = sum(area_m2, na.rm = TRUE), avg_build_size = mean(area_m2, na.rm = TRUE)), by = grid_id]
  dplyr::left_join(grid_sf, as.data.frame(grid_stats), by = "grid_id") |>
    dplyr::mutate(
      n_buildings    = dplyr::if_else(is.na(n_buildings), 0L, n_buildings),
      total_area_m2  = dplyr::if_else(is.na(total_area_m2), 0,  total_area_m2),
      avg_build_size = dplyr::if_else(is.na(avg_build_size), 0,  avg_build_size)
    )
}

# Subplace rollup by **buildings (centroid)** 
rollup_subplaces_by_building <- function(buildings_sa, admin_keep, out_neigh_gpkg, out_neigh_csv) {
  sp_all <- sf::st_read(SUBPLACE_FILE, quiet = TRUE) |> sf::st_make_valid()
  if (is.na(sf::st_crs(sp_all))) sf::st_crs(sp_all) <- CRS_WGS84
  if (!is.null(admin_keep) && "DC_NAME" %in% names(sp_all)) {
    city_sp <- sp_all[sp_all$DC_NAME %in% admin_keep, ]
    if (nrow(city_sp) == 0) stop("No subplaces matched DC_NAME. Check admin_keep.")
  } else {
    city_sp <- sp_all
  }
  city_sp <- city_sp |>
    sf::st_transform(sf::st_crs(buildings_sa)) |>
    sf::st_make_valid() |>
    sf::st_set_precision(0.01)
  # Clean to polygonal only
  city_sp <- sf::st_collection_extract(city_sp, "POLYGON", warn = FALSE)
  city_sp <- city_sp[!sf::st_is_empty(city_sp), , drop = FALSE]
  city_sp <- sf::st_buffer(city_sp, 0)
  stopifnot("SP_CODE" %in% names(city_sp))
  
  # Building centroid → subplace join (one-to-one)
  b_pts <- sf::st_centroid(buildings_sa)
  j <- sf::st_join(
    sf::st_sf(bid = seq_len(nrow(b_pts)), geometry = sf::st_geometry(b_pts)),
    city_sp[, c("SP_CODE")], left = FALSE
  )
  if (!nrow(j)) {
    warning("No buildings fell into subplaces.")
    return(invisible(FALSE))
  }
  dt <- data.table::data.table(
    SP_CODE = as.character(j$SP_CODE),
    area_m2 = buildings_sa$area_m2[j$bid]
  )
  neigh_stats <- dt[, .(
    n_buildings_total = .N,
    total_area_m2     = sum(area_m2, na.rm = TRUE)
  ), by = SP_CODE][
    , avg_building_size := fifelse(n_buildings_total > 0, total_area_m2 / n_buildings_total, 0)
  ]
  
  # Join back to polygons for output
  city_sp$SP_CODE <- as.character(city_sp$SP_CODE)
  city_sp_out <- dplyr::left_join(city_sp, as.data.frame(neigh_stats), by = "SP_CODE") |>
    dplyr::mutate(dplyr::across(where(is.numeric), ~ ifelse(is.na(.), 0, .)))
  
  st_write(city_sp_out, out_neigh_gpkg, delete_dsn = TRUE, quiet = TRUE)
  data.table::fwrite(
    st_drop_geometry(
      city_sp_out[, c(intersect(c("SP_CODE","SP_NAME","MP_NAME","DC_NAME"), names(city_sp_out)),
                      "n_buildings_total","total_area_m2","avg_building_size")]
    ),
    out_neigh_csv
  )
  
  invisible(TRUE)
}

## --------- Core runner per city ----------
process_city <- function(city_id, parquet_file, boundary_file, admin_keep) {
  
  stopifnot(file.exists(parquet_file), file.exists(boundary_file))
  
  # Outputs
  OUT_BUILDINGS_WGS84_GPKG <- file.path(OUT_DIR, sprintf("%s_buildings_wgs84_clip.gpkg", city_id))
  OUT_BUILDINGS_SA_GPKG    <- file.path(OUT_DIR, sprintf("%s_buildings_sa_albers.gpkg", city_id))
  OUT_GRID_GPKG            <- file.path(OUT_DIR, sprintf("%s_grid_200m_building_stats.gpkg", city_id))
  OUT_NEIGH_GPKG           <- file.path(OUT_DIR, sprintf("%s_subplaces_building_stats.gpkg", city_id))
  OUT_GRID_CSV             <- file.path(OUT_DIR, sprintf("%s_grid_200m_building_stats.csv", city_id))
  OUT_NEIGH_CSV            <- file.path(OUT_DIR, sprintf("%s_subplaces_building_stats.csv", city_id))
  TMP_DIR <- file.path(OUT_DIR, sprintf("tmp_chunks_%s", city_id))
  if (dir.exists(TMP_DIR)) unlink(TMP_DIR, recursive = TRUE, force = TRUE)
  dir.create(TMP_DIR, recursive = TRUE, showWarnings = FALSE)
  for (p in c(OUT_BUILDINGS_WGS84_GPKG, OUT_BUILDINGS_SA_GPKG, OUT_GRID_GPKG, OUT_NEIGH_GPKG, OUT_GRID_CSV, OUT_NEIGH_CSV)) {
    if (file.exists(p)) file.remove(p)
  }
  
  # Boundary (from prebuilt GPKG)
  B <- build_city_boundary(admin_keep, boundary_file)
  BOUND_WGS84      <- B$BOUND_WGS84
  BOUND_WGS84_GEOM <- B$BOUND_WGS84_GEOM
  BOUND_BBOX_SFC   <- B$BOUND_BBOX_SFC
  BOUND_SA         <- B$BOUND_SA
  
  # Sanity: area + expected grid count
  area_km2  <- as.numeric(units::set_units(sf::st_area(BOUND_SA), km^2))
  est_grids <- round((area_km2 * 1e6) / (CELL_M^2))
  
  
  # Read GeoParquet → ensure WGS84
  b_sf <- sfarrow::st_read_parquet(parquet_file)
  if (is.na(st_crs(b_sf))) st_crs(b_sf) <- CRS_WGS84
  if (st_crs(b_sf)$epsg != 4326) b_sf <- st_transform(b_sf, CRS_WGS84)
  
  # Candidate bbox filter in WGS84 (s2 on)
  old_s2 <- sf::sf_use_s2(); on.exit(sf::sf_use_s2(old_s2), add = TRUE)
  sf::sf_use_s2(TRUE)
  idx0 <- st_intersects(b_sf, BOUND_BBOX_SFC, sparse = TRUE, prepared = TRUE)
  b_sf <- b_sf[lengths(idx0) > 0L, , drop = FALSE]
  n <- nrow(b_sf)
  if (n == 0) { warning("No buildings in bbox for ", city_id); return(invisible(FALSE)) }
  
  
  # Chunk + WGS84 boundary clip → temp files
  starts <- seq(1L, n, by = BATCH_SIZE)
  ends   <- pmin(starts + BATCH_SIZE - 1L, n)
  slices <- Map(seq.int, starts, ends)
  
  process_chunk <- make_process_chunk_wgs84(BOUND_BBOX_SFC, BOUND_WGS84_GEOM, TMP_DIR)
  chunk_files <- future_map2_chr(
    .x = seq_along(slices),
    .y = slices,
    .f = function(i, ix) process_chunk(i, b_sf[ix, , drop = FALSE]),
    .options = furrr_options(seed = TRUE, scheduling = 1)
  )
  rm(b_sf); invisible(gc())
  chunk_files <- chunk_files[!is.na(chunk_files)]
  if (!length(chunk_files)) { warning("No intersected buildings written for ", city_id); return(invisible(FALSE)) }
  
  
  # Merge WGS84 chunks → consolidated WGS84 buildings
  merge_chunk_files_wgs84(chunk_files, OUT_BUILDINGS_WGS84_GPKG)
  
  
  # Project consolidated to SA-Albers and compute area_m2
  buildings_sa <- st_read(OUT_BUILDINGS_WGS84_GPKG, quiet = TRUE)
  if (st_crs(buildings_sa) != CRS_SA_ALBERS) buildings_sa <- st_transform(buildings_sa, CRS_SA_ALBERS)
  buildings_sa <- st_make_valid(buildings_sa)
  buildings_sa <- st_set_precision(buildings_sa, 0.01)
  if (SIMPLIFY_BUILDINGS_M > 0) buildings_sa <- st_simplify(buildings_sa, dTolerance = SIMPLIFY_BUILDINGS_M)
  buildings_sa$area_m2 <- as.numeric(drop_units(st_area(buildings_sa)))
  
  
  if (WRITE_BUILDINGS_GPKG) {
    st_write(buildings_sa, OUT_BUILDINGS_SA_GPKG, delete_dsn = TRUE, quiet = TRUE)
  }
  
  # 200 m grid within SA-Albers boundary
  grid <- st_make_grid(BOUND_SA, cellsize = CELL_M, square = TRUE)
  if (length(grid) == 0) stop("Boundary produced zero grid cells. Check boundary CRS/source.")
  grid_sf <- st_sf(grid_id = seq_along(grid), geometry = grid) |>
    st_intersection(BOUND_SA) |>
    st_make_valid() |>
    st_set_precision(0.01)
  
  
  # Grid stats
  grid_sf <- grid_stats_from_buildings(buildings_sa, grid_sf)
  
  
  # Write grid outputs
  st_write(grid_sf, OUT_GRID_GPKG, delete_dsn = TRUE, quiet = TRUE)
  data.table::fwrite(st_drop_geometry(grid_sf), OUT_GRID_CSV)
  
  # Subplace rollup by buildings (no duplication)
  rollup_subplaces_by_building(
    buildings_sa   = buildings_sa,
    admin_keep     = admin_keep,
    out_neigh_gpkg = OUT_NEIGH_GPKG,
    out_neigh_csv  = OUT_NEIGH_CSV
  )
  
  # Cleanup
  try(unlink(TMP_DIR, recursive = TRUE, force = TRUE), silent = TRUE)
  
  invisible(TRUE)
}

## --------- Driver ----------
invisible(process_city(
  city_id      = city_cfg$capetown$city_id,
  parquet_file = city_cfg$capetown$parquet_file,
  boundary_file= city_cfg$capetown$boundary_file,
  admin_keep   = city_cfg$capetown$admin_keep
))

invisible(process_city(
  city_id      = city_cfg$johannesburg$city_id,
  parquet_file = city_cfg$johannesburg$parquet_file,
  boundary_file= city_cfg$johannesburg$boundary_file,
  admin_keep   = city_cfg$johannesburg$admin_keep
))

invisible(process_city(
  city_id      = city_cfg$durban$city_id,
  parquet_file = city_cfg$durban$parquet_file,
  boundary_file= city_cfg$durban$boundary_file,
  admin_keep   = city_cfg$durban$admin_keep
))



